from concurrent.futures import ProcessPoolExecutor
from concurrent.futures import ThreadPoolExecutor
from sentinelsat import SentinelAPI
from queue import Queue
import pandas as pd
import datetime
import paramiko
import getpass
import argparse
# import shutil
import time
import json
# import math
import os
# import re

# 获取日期列表
def get_dates(start_date, end_date):
	""""""
	year, month, day = start_date
	start_date = datetime.datetime(year, month, day)
	year, month, day = end_date
	end_date = datetime.datetime(year, month, day)
	dates = []
	while start_date < end_date:
		dates.append(start_date)
		start_date += datetime.timedelta(days=+1)
	return dates

# 进度条
def progress_bar(start_time, index, tasks, bar_long=40):
	""""""
	total = len(tasks)
	index = index + 1
	percent = index / total
	now_time = time.time()
	print(f"{'█' * round(percent * bar_long)}", end='')
	print(f"{' ' * round(bar_long - percent * bar_long)}", end='')
	print(f":{round(percent * 100)}%({index}/{total})", end='')
	estimated_time = round((now_time-start_time)/index*total)
	print(f"(预计耗时:{estimated_time//60}m:{estimated_time%60}s", end='')
	actual_time = round((now_time-start_time))
	print(f",实际耗时:{actual_time//60}m:{actual_time%60}s", end='')
	remaining_time = estimated_time-actual_time
	print(f",剩余时间:{remaining_time//60}m:{remaining_time%60}s", end='')
	speed = round(actual_time / index)
	print(f",实际速度:{speed}s/个)", end='\r')
	return

# 检索esa数据
def search_l1c_zip_count_xlsx(
		username,
		password,
		tiles,
		dates,
		maxcloud,
		xlsx_path,
		json_path=None):
	""""""
	if isinstance(dates, str):
		start_date = datetime.datetime.strptime(dates, "%Y%m")
		end_date = datetime.datetime(
			start_date.year + (start_date.month == 12),
			start_date.month * (start_date.month != 12) + 1, 1)
		start_date = [start_date.year, start_date.month, start_date.day]
		end_date = [end_date.year, end_date.month, end_date.day]
		dates = get_dates(start_date, end_date)
	if isinstance(tiles, str):
		province = tiles

		if json_path is None:
			json_path = r"/home/冯洋洋/PythonProjection/RasterShape/conf/ChinaMgrs.json"
		with open(json_path, "r", encoding="UTF-8") as jsonFile:
			china_tiles = json.load(jsonFile)
		tiles = china_tiles[province]

		sheet_name = f"{province}_{dates[0].strftime('%Y%m%d')}-{dates[-1].strftime('%Y%m%d')}"
	else:
		sheet_name = f"{dates[0].strftime('%Y%m%d')}-{dates[-1].strftime('%Y%m%d')}"
	api = SentinelAPI(username, password, timeout=1)
	query_kwargs = {
		'platformname': 'Sentinel-2',
		'producttype': 'S2MSI1C',
		'date': (dates[0], dates[-1]+datetime.timedelta(days=1)),
		"cloudcoverpercentage": (0, maxcloud)}
	product_count_dataframe = pd.DataFrame(
		columns=[date.strftime("%Y%m%d") for date in dates], index=tiles)
	start_time = time.time()
	for index, tile in enumerate(tiles):
		query_kwargs["raw"] = f'tileid:{tile} OR filename:*_T{tile}_*'  # 防止离线数据下载失败
		query_kwargs["tileid"] = tile
		products = api.query(**query_kwargs)
		for product in products:
			product_info = api.get_product_odata(product, full=True)
			product_date = product_info['date'].strftime('%Y%m%d')
			if product_count_dataframe[product_date][tile] >= 1:
				product_count_dataframe[product_date][tile] += 1
			else:
				product_count_dataframe[product_date][tile] = 1
		progress_bar(start_time, index, tiles)
	if os.path.exists(xlsx_path):
		with pd.ExcelWriter(xlsx_path, mode="a", engine='openpyxl') as writer:
			product_count_dataframe.to_excel(writer, sheet_name=sheet_name)
	else:
		with pd.ExcelWriter(xlsx_path, mode="w", engine='openpyxl') as writer:
			product_count_dataframe.to_excel(writer, sheet_name=sheet_name)
	return

# 计算文件变化速度, 单位Mb/s
def get_file_change_speed(file_path, max_time):
	""""""
	source_file_size = os.path.getsize(file_path)
	time.sleep(max_time)
	target_file_size = os.path.getsize(file_path)
	change_size = (target_file_size - source_file_size)/(1024*1024)/max_time
	print(change_size)
	return change_size


def download_one_sentinel2_L1c(
		username, password, tile, date, save_folder_path, maxcloud=100):
	try:
		api = SentinelAPI(username, password, timeout=60)
		query_kwargs = {
			'platformname': 'Sentinel-2',
			'producttype': 'S2MSI1C',
			'date': (date, date + datetime.timedelta(days=1)),
			"cloudcoverpercentage": (0, maxcloud),
			"tileid": tile
		}
		products = api.query(**query_kwargs)
		for product in products:
			if api.is_online(product):
				api.download(product, directory_path=save_folder_path)
		return True
	except:
		print(f"{tile},{date.strftime('%Y%m%d')},下载失败,重新下载!")
		return False

# 
def sentinel2_l1c_download(
		username, password, tile, date, save_folder_path, maxcloud=100, attempt=5):
	""""""
	count = 1
	while count <= attempt:
		if count != 1:
			print(f"第{count}次尝试下载!")
		flag = download_one_sentinel2_L1c(
			username, password, tile, date, save_folder_path, maxcloud)
		if flag:
			return
		else:
			time.sleep(2)
			count += 1
	return

# 按照图幅列表与时期下载哨兵数据
def download_sentinel2_L1c(
		username, password, tiles, dates, maxcloud, save_folder_path):
	""""""
	api = SentinelAPI(username,  password, timeout=30)
	query_kwargs = {
		'platformname': 'Sentinel-2',
		'producttype': 'S2MSI1C',
		'date': (dates[0], dates[-1]+datetime.timedelta(days=1)),
		"cloudcoverpercentage": (0, maxcloud)}
	for tile in tiles:
		time.sleep(1)
		query_kwargs["raw"] = f'tileid:{tile} OR filename:*_T{tile}_*'  # 防止离线数据下载失败
		query_kwargs["tileid"] = tile  # 单独直接以图幅号搜索存在缺陷, 离线数据会下载失败。
		product_list = api.query(**query_kwargs)
		for product in product_list:
			product_info = api.get_product_odata(product, full=True)
			print(product_info['title'])
			api.download(product, directory_path=save_folder_path)
	print(f"{username}用户任务完成!")
	return


def search_esa_sentinel2_l1c(
		username, password, tile, yearmonth, maxcloud=100):
	""""""
	start_date = datetime.datetime.strptime(yearmonth, "%Y%m")
	end_date = datetime.datetime(
		start_date.year + (start_date.month == 12),
		start_date.month * (start_date.month != 12) + 1,
		1
	)
	api = SentinelAPI(username, password, timeout=30)
	query_kwargs = {'platformname': 'Sentinel-2',
					'producttype': 'S2MSI1C',
					'date': (start_date, end_date),
					"cloudcoverpercentage": (0, maxcloud),
					"raw": f'tileid:{tile} OR filename:*_T{tile}_*',
					"tileid": tile}
	products = [product for product in api.query(**query_kwargs)]
	return products


def scp_file(zip_path, date_dir):
	command = "scp {zip_path} root@192.168.2.172:{date_dir}"
	os.system(command)
	return


def sub_thread_download_esa_sentinel2_l1c(
		product, 
		users, 
		userq, 
		save_dir,
		remote_dir=r"/home/遥感技术部/遥感数据/Sentinel2L1cZip", 
		timeout=180):
	user_index = userq.get(block=True, timeout=1)
	[username, password] = users[user_index]
	
	try:
		ssh = paramiko.SSHClient()
		ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
		ssh.connect('192.168.2.172', 22, "root", "DELL123dell")
		sftp = ssh.open_sftp()
		api = SentinelAPI(username, password, timeout=timeout)
		product_info = api.get_product_odata(product, full=True)
		product_date = product_info["date"]
		product_name = product_info['title']
		date_name = product_date.strftime("%Y%m%d")

		zip_name = f"{product_name}.zip"
		remote_path = "/".join([remote_dir, date_name, zip_name])
		try:
			sftp.stat(remote_path)
			print(f"{datetime.datetime.now()}:{zip_name}文件存在!")
		except:
			if api.is_online(product):
				try:
					api.download(product, directory_path=save_dir)
					local_path = os.path.join(save_dir, zip_name)
					date_dir = "/".join([remote_dir, date_name])
					try:
						sftp.stat(date_dir)
					except:
						ssh.exec_command(f"mkdir {date_dir}")
					print(f"{product_name}:下载成功!\n")
					sftp.put(local_path, remote_path)
					os.remove(local_path)			
				except:
					print(f"{product_name}>>>下载失败!\n")		
					# os.remove(os.path.join(save_dir, zip_name+".incomplete"))	
			else:
				print(f"{product_name}:离线数据!\n")
	except:
		print(f"{datetime.datetime.now()}:>因网络问题下载失败!")
	ssh.close()
	userq.put(user_index, block=True, timeout=1)
	return 		


def sub_thread_search_esa_sentinel2_l1c(
		tile, users, userq, yearmonth, maxcloud=100, attempt=5):
	""""""
	user_index = userq.get(block=True, timeout=1)
	[username, password] = users[user_index]
	start_date = datetime.datetime.strptime(yearmonth, "%Y%m")
	end_date = datetime.datetime(
		start_date.year + (start_date.month == 12),
		start_date.month * (start_date.month != 12) + 1,
		1
	)
	products = []
	while attempt > 0:
		try:
			api = SentinelAPI(username, password, timeout=30)
			raw = f'tileid:{tile} OR filename:*_T{tile}_*'
			query_kwargs = {'platformname': 'Sentinel-2',
							'producttype': 'S2MSI1C',
							'date': (start_date, end_date),
							"cloudcoverpercentage": (0, maxcloud),
							"raw": raw,
							"tileid": tile}
			products = [
				product for product in api.query(**query_kwargs)
				]
			print(f"{username}:{tile}:{len(products)}")
			userq.put(user_index, block=True, timeout=1)
			return products
		except:
			attempt -= 1
	userq.put(user_index, block=True, timeout=1)
	print(f"{username}:{tile}:{len(products)}")
	return products

def muitl_thread_download_l1c(
		users,
		tiles,
		yearmonth=None,
		save_folder_path='',
		maxcloud=100):
	""""""
	if yearmonth is None:
		now_date = datetime.datetime.now()
		yearmonth = now_date.strftime("%Y%m")
	userq = Queue(maxsize=len(users))
	for i in range(len(users)):
		userq.put(i)
	# 多线程查询esa发布的数据
	with ThreadPoolExecutor(max_workers=len(users)) as executor:
		results = executor.map(
			sub_thread_search_esa_sentinel2_l1c,
			tiles,
			[users] * len(tiles),
			[userq] * len(tiles),
			[yearmonth] * len(tiles),
			[maxcloud] * len(tiles)
		)
	proucts = [prouct for result in results for prouct in result]
	print(f"共计{len(proucts)}景数据")
	# 多线程下载esa发布的数据
	with ThreadPoolExecutor(max_workers=len(users)) as executor:
		results = executor.map(
			sub_thread_download_esa_sentinel2_l1c,
			proucts,
			[users]*len(proucts),
			[userq]*len(proucts),
			[save_folder_path]*len(proucts)
		)
	# 多线程下载失败的数据
	for prouct in results:
		if prouct is not None:
			print(f"{prouct}:下载失败!")
	print("END!")
	return

# get sentinel2 tile list
def get_tiles(
    province=None, 
    city=None, 
    county=None, 
    tiles=[]):
    con = psycopg2.connect(
    database="sentinel2",
    user="postgres",
    password="12345",
    host="192.168.2.172",
    port="5432"
    )
    cur = con.cursor()
    if province is not None:
        sql = f"""
        SELECT "中国有效哨兵格网分布wgs84".mgrs
        FROM china_sy
        JOIN "中国有效哨兵格网分布wgs84" 
        on st_intersects(china_sy.geom,"中国有效哨兵格网分布wgs84".geom)
        WHERE  china_sy.qhmc = '{province}' 
        """
    elif city is not None:
        sql = f"""
        SELECT "中国有效哨兵格网分布wgs84".mgrs
        FROM china_dsy
        JOIN "中国有效哨兵格网分布wgs84" 
        on st_intersects(china_dsy.geom,"中国有效哨兵格网分布wgs84".geom)
        WHERE  china_dsy.qhmc = '{city}' 
        """
    elif county is not None:
        sql = f"""
        SELECT "中国有效哨兵格网分布wgs84".mgrs
        FROM china_xy
        JOIN "中国有效哨兵格网分布wgs84" 
        on st_intersects(china_xy.geom,"中国有效哨兵格网分布wgs84".geom)
        WHERE  china_xy.qhmc = '{county}' 
        """
    cur.execute(sql)
    con.commit()
    tiles.extend([record[0] for record in cur.fetchall()])
    tiles = [tile.upper() for tile in tiles]
    cur.close()
    con.close()
    return tiles


def main():
	""""""
	'''多线程 esa 下载数据'''
	# provinces = ["河北省", "山东省", "河南省", "山西省", "江苏省", "东四盟", "甘肃省"]
	# provinces = ["河北省", "山东省", "河南省", "山西省", "江苏省"]
	provinces = ["江苏省"]
	yearmonth = "202203"
	save_folder_path = "E:\\Yang\\Project\\PrivateToolBox\\data\\temp"
	# yearmonth = None
	users = [
		["fyy001", "show1234"],
		["fyy_01", "show1234"],
		["fyy_02", "show1234"],
		["fyy_03", "show1234"],
		["fyy_04", "show1234"],
		["hxx2021", "hxx2021123"],
		["hxx2022", "hxx2022123"],
		["hxx20221", "hxx2022123"],
		["hxx20222", "hxx20222123"],
		["wyq01", "show970213"],
		["wyq02", "show970213"],
		["wyq03", "show970213"],
		["wyq04", "show970213"],
		["wyq05", "show970213"],
		["wyq06", "show970213"],
		["wyq07", "show970213"],
	]
	# json_path = "E:\\Yang\\Project\\PrivateToolBox\\conf\\ChinaMgrs.json"
	# 获得该多个省的图幅列表
	# with open(json_path, "r", encoding="UTF-8") as jsonFile:
	# 	china_tiles = json.load(jsonFile)
	tiles = list(set([
			tile 
			for province in provinces 
			for tile in get_tiles(province=province)
			]))
	
	while True:
		muitl_thread_download_l1c(
			users, 
			tiles, 
			yearmonth, 
			save_folder_path)
		print("END!")
		time.sleep(60)


if __name__ == "__main__":
	main()